#include "CCoroutineImpl.h"
#include "CCoroutine.h"
#include <exception>
#include "../promise.h"
#include "CCoroutineTimerContext.h"
#include "CCoroutineGroupImpl.h"

namespace TB
{
	std::once_flag CCoroutineImpl::m_oInitFlag;

	CMPSCQueue<std::pair<UINT64, std::weak_ptr<CCoroutineImpl>>> CCoroutineImpl::m_oSetTimeoutQueue;

	std::shared_ptr<CTimerMgr> CCoroutineImpl::m_poTimeoutTimeMgr;

	std::atomic_uint64_t CCoroutineImpl::m_qwIncID;

	CMPSCQueue<STimerUID> CCoroutineImpl::m_oDestoryTimerQueue;

	CCoroutineImpl::CCoroutineImpl(const std::weak_ptr<CCoroutineGroupImpl>& poGroup, const std::weak_ptr<CThreadPool>& poThreadPool, DSetCurCoroutine cbSetCurCoroutine, bool bManualPoll):
		m_poGroup(poGroup), m_poThreadPool(poThreadPool), m_cbSetCurCoroutine(cbSetCurCoroutine), 
		m_qwID(m_qwIncID.fetch_add(1, std::memory_order_relaxed) + 1), m_bManualPoll(bManualPoll)
	{
		if (m_qwID == UINT64_MAX)
			throw new std::runtime_error("coroutine's id limit");
		if (!m_bManualPoll)
		{
			std::call_once(m_oInitFlag, []()
			{
				CThreadPool::PollPerMilliSecond(std::bind(&CCoroutineImpl::_PollPerMilliSecondInThreadPool));
				m_poTimeoutTimeMgr = std::make_shared<CTimerMgr>(false, CThreadPool::GetCurTimePtr());
			});
		}
		else
		{
			m_oLastPollTime = std::chrono::steady_clock::now();
		}
		m_poSyncContext = std::make_shared<CMPSCSyncContext>();
		m_poTimerContext = std::make_shared<CCoroutineTimerContext>(m_bManualPoll ? nullptr : CThreadPool::GetCurTimePtr());
		if (m_poThreadPool.expired() && !m_bManualPoll)
			m_poThreadPool = CThreadPool::Default();
	}

	CCoroutineImpl::~CCoroutineImpl()
	{
		if (m_sWakeupTimerUID)
			m_oDestoryTimerQueue.Push(m_sWakeupTimerUID);
	}

	void CCoroutineImpl::Post(const std::function<void()>* pcbFunction, bool bFromGroup)
	{
		if (pcbFunction != nullptr)
			m_poSyncContext->Post(*pcbFunction);
		if (!m_bManualPoll && CCoroutine::GetCurrent() != m_poCoroutine.lock())
		{
			auto dwSem = m_oSem.fetch_add(1, std::memory_order_relaxed);
			if (dwSem < 0) [[unlikely]]
				throw std::runtime_error("sem < 0");
			else if (dwSem == 0)
			{
				auto poImpl = m_poSelf.lock();
				if (poImpl)
				{
					auto poThreadPool = poImpl->m_poThreadPool.lock();
					if (poThreadPool)
					{
						poThreadPool->Run([poImpl, bFromGroup]()
						{
							poImpl->Poll(bFromGroup);
						});
					}
				}
			}
		}
	}

	UINT32 CCoroutineImpl::Poll(bool bFromGroup)
	{
		auto poCoroutine = m_poCoroutine.lock();
		if (!poCoroutine)
			return 0;
		std::shared_ptr<CCoroutineGroupImpl> poGroup = m_poGroup.lock();
		if (m_oPollLock.fetch_add(1, std::memory_order_relaxed) != 0)
		{
			if (bFromGroup)
				poGroup->m_oIdleCoroutineQueue.Push(m_poSelf.lock());
			return 0;
		}

		UINT32 dwCount = 0;

		do
		{
			if (poGroup)
				poGroup->m_oIdleCount.fetch_sub(1, std::memory_order_relaxed);

			// set context
			auto poPrevCoroutine = CCoroutine::GetCurrent();
			m_cbSetCurCoroutine(poCoroutine);
			CCurThreadSyncContext& roCurThreadSyncContext = CCurThreadSyncContext::GetCurrent();
			auto poPrevSyncContext = roCurThreadSyncContext.GetSynchronizationContext();
			roCurThreadSyncContext.SetSynchronizationContext(m_poSyncContext);
			CCurThreadTimerContext& roCurThreadTimerContext = CCurThreadTimerContext::GetCurrent();
			auto poPrevTimerContext = roCurThreadTimerContext.GetTimerContext();
			roCurThreadTimerContext.SetTimerContext(m_poTimerContext);

			CTimerMgr& roTimerMgr = m_poTimerContext->GetTimerMgr();

			dwCount += _Poll();

			// reset context
			roCurThreadTimerContext.SetTimerContext(poPrevTimerContext);
			roCurThreadSyncContext.SetSynchronizationContext(poPrevSyncContext);
			m_cbSetCurCoroutine(poPrevCoroutine);

			// yield
			if (!m_bManualPoll && m_oSem.load(std::memory_order_acquire) == 0)
			{
				UINT64 qwNextPollTime = roTimerMgr.GetNextPollTime();
				if (qwNextPollTime != (UINT64)-1 && (m_qwWakeupTime == 0 || m_qwWakeupTime != qwNextPollTime))
				{
					_WakeupOn(qwNextPollTime);
				}
			}

			if (poGroup)
			{
				UINT32 dwCoroutineCount = poGroup->GetSize();
				UINT32 dwIdleCount = poGroup->m_oIdleCount.fetch_add(1, std::memory_order_release);
				if (dwIdleCount != dwCoroutineCount && dwCoroutineCount <= dwIdleCount + poGroup->m_oTaskQueue.GetCount())
					continue;
			}
			if (m_oPollLock.fetch_sub(1, std::memory_order_release) == 1)
				break;
		} while (true);

		if (bFromGroup)
		{
			poGroup->m_oIdleCoroutineQueue.Push(m_poSelf.lock());
		}

		return dwCount;
	}

	UINT32 CCoroutineImpl::_Poll()
	{
		UINT32 dwCount = 0;
		UINT32 dwPollTimerCount = 0, dwPollAsyncTaskCount = 0, dwPollMpmcAsyncTaskCount = 0;
		CTimerMgr& roTimerMgr = m_poTimerContext->GetTimerMgr();
		CMPSCSyncContext& roMPSCSyncContext = *m_poSyncContext;
		std::shared_ptr<CCoroutineGroupImpl> poGroup = m_poGroup.lock();

		UINT32 dwCount1 = 0;
		do
		{
			dwCount1 = 0;
			// timers
			UINT32 dwMaxCountPerPoll = roTimerMgr.GetMaxCountPerPoll();
			UINT64 qwTimerPollDeltaTime = 0;
			if (m_bManualPoll)
			{
				auto oLastPollTime = m_oLastPollTime;
				m_oLastPollTime = std::chrono::steady_clock::now();
				qwTimerPollDeltaTime = (UINT64)((m_oLastPollTime - oLastPollTime).count() / 1000000);
			}
			dwPollTimerCount = roTimerMgr.Poll(qwTimerPollDeltaTime);
			dwCount1 += dwPollTimerCount;

			// mpsc tasks
			INT32 dwSem = 0;
			dwPollAsyncTaskCount = 0;
			INT32 dwTmpCount = 0;
			do
			{
				dwSem = m_oSem.load(std::memory_order_acquire);
				dwTmpCount = 0;
				while ((dwMaxCountPerPoll == 0 || dwPollAsyncTaskCount < dwMaxCountPerPoll) && !roMPSCSyncContext.IsEmpty())
				{
					++dwTmpCount;
					++dwPollAsyncTaskCount;
					const std::function<void()>& rcbFunction = roMPSCSyncContext.Front();
					rcbFunction();
					roMPSCSyncContext.Pop();
				}
			} while ((dwMaxCountPerPoll == 0 || dwPollAsyncTaskCount < dwMaxCountPerPoll) && (dwTmpCount != 0 || (dwSem != 0 && m_oSem.fetch_sub(dwSem, std::memory_order_relaxed) != dwSem)));
			dwCount1 += dwPollAsyncTaskCount;

			// mpmc tasks
			if (poGroup)
			{
				dwPollMpmcAsyncTaskCount = 0;
				std::function<void()> cbFunc;
				while ((dwMaxCountPerPoll == 0 || dwPollMpmcAsyncTaskCount < dwMaxCountPerPoll) && 0 != poGroup->m_oTaskQueue.Pop(cbFunc))
				{
					++dwPollMpmcAsyncTaskCount;
					cbFunc();
				}
				dwCount1 += dwPollMpmcAsyncTaskCount;
			}
			dwCount += dwCount1;
		} while (dwCount1 != 0);

		return dwCount;
	}
	
	void CCoroutineImpl::_PollPerMilliSecondInThreadPool()
	{
		CTimerMgr& roTimeoutTimerMgr = *m_poTimeoutTimeMgr;
		while (!m_oSetTimeoutQueue.IsEmpty())
		{
			const std::pair<UINT64, std::weak_ptr<CCoroutineImpl>>& roPair = m_oSetTimeoutQueue.Front();
			std::shared_ptr<CCoroutineImpl> poPtr = roPair.second.lock();
			if (poPtr != nullptr)
			{
				if (poPtr->m_sWakeupTimerUID)
				{
					roTimeoutTimerMgr.StopTimer(poPtr->m_sWakeupTimerUID);
					poPtr->m_sWakeupTimerUID = {};
				}
				UINT64 qwCurTime = roTimeoutTimerMgr.GetCurTime();
				UINT64 qwDelta = qwCurTime >= roPair.first ? 0 : roPair.first - qwCurTime;
				if (qwDelta == 0)
					_Wakeup(roPair.first, roPair.second);
				else
					poPtr->m_sWakeupTimerUID = roTimeoutTimerMgr.StartTimer(qwDelta, 1, std::bind(&CCoroutineImpl::_Wakeup, roPair.first, roPair.second));
			}
			m_oSetTimeoutQueue.Pop();
		}
		for (;!m_oDestoryTimerQueue.IsEmpty(); m_oDestoryTimerQueue.Pop())
		{
			roTimeoutTimerMgr.StopTimer(m_oDestoryTimerQueue.Front());
		}
		roTimeoutTimerMgr.Poll(0);
	}

	void CCoroutineImpl::_Wakeup(UINT64 qwWakeupTime, const std::weak_ptr<CCoroutineImpl>& poWeakPtr)
	{
		std::shared_ptr<CCoroutineImpl> poImpl = poWeakPtr.lock();
		if (poImpl)
		{
			auto poThreadPool = poImpl->m_poThreadPool.lock();
			if (poThreadPool)
			{
				poThreadPool->Run([qwWakeupTime, poImpl]()
				{
					if (poImpl->m_qwWakeupTime == qwWakeupTime)
						poImpl->m_qwWakeupTime = 0;
					poImpl->Poll(false);
				});
			}
		}
	}

	void CCoroutineImpl::_WakeupOn(UINT64 qwWakeupTime)
	{
		m_qwWakeupTime = qwWakeupTime;
		m_oSetTimeoutQueue.Push(std::make_pair(m_qwWakeupTime, m_poSelf));
	}
}